package com.yujiahao;

import org.apache.flink.kinesis.shaded.com.amazonaws.ClientConfiguration;
import org.apache.flink.kinesis.shaded.com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
import org.apache.flink.kinesis.shaded.com.amazonaws.services.kinesis.AmazonKinesis;
import org.apache.flink.kinesis.shaded.com.amazonaws.services.kinesis.AmazonKinesisClientBuilder;
import org.apache.flink.kinesis.shaded.com.amazonaws.services.kinesis.model.*;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import java.nio.charset.StandardCharsets;
import java.util.List;

/**
 * @Description:  使用分片迭代器
 * @Author: yujiahao
 * @Date: 2023/8/1 15:08
 * @Version: 1.0
 */

public class KinesisConsume {
    private static AmazonKinesis client;
    private String streamName = "test-kds-api";
    DescribeStreamRequest describeStreamRequest;
    /**
     * 创建客户端对象
     */
    @Before
    public void createClient() {
        //构建客户端对象
        AmazonKinesisClientBuilder clientBuilder = AmazonKinesisClientBuilder.standard();
        ClientConfiguration clientConfiguration = new ClientConfiguration();
        clientBuilder.setCredentials(new DefaultAWSCredentialsProviderChain());
        clientBuilder.setClientConfiguration(clientConfiguration);
        client = clientBuilder.build();
        describeStreamRequest = new DescribeStreamRequest().withStreamName(streamName);
    }


    /**
     * 获取数据流中所有分片的数据并打印
     */
    @Test
    public void getData() {
        // 获取数据流的描述信息
        DescribeStreamResult describeStreamResult = client.describeStream(streamName);
        List<Shard> shards = describeStreamResult.getStreamDescription().getShards();
        // 读数据还是要从分片中读取的，这里是所有活跃的，当然也可以指定读取
        for (Shard shard : shards) {
            // 判断分片是否处于活跃状态,不判断的话可能会把历史不可用的分区都扫描一遍，效率极差
            SequenceNumberRange range = shard.getSequenceNumberRange();
            //如果分片的结束序列号为空，意味着该分片仍然处于活跃状态
            if (range.getEndingSequenceNumber() == null) {
                // 获取分片的迭代器
                String shardIterator = getShardIterator(shard.getShardId());
                //这里是分片的状态，不为空说明有数
                while (shardIterator != null) {
                    // 创建获取数据的请求
                    GetRecordsRequest request = new GetRecordsRequest();
                    request.setShardIterator(shardIterator);
                    request.setLimit(1000); // 设置每次获取的最大记录数，可根据需求调整

                    // 获取数据记录
                    GetRecordsResult result = client.getRecords(request);
                    List<Record> records = result.getRecords();

                    // 处理数据，这里只是简单打印每条记录的数据
                    for (Record record : records) {
                        String data = new String(record.getData().array(), StandardCharsets.UTF_8);
                        System.out.println("Received data: " + data);
                    }

                    // 获取下一批数据的迭代器
                    shardIterator = result.getNextShardIterator();
                }
            }
        }
    }

    /**
     * 获取分片的ShardIterator，从最早记录（历史数据）开始读取
     *
     * @param shardId 分片ID
     * @return 分片的ShardIterator
     *
     * todo withShardIteratorType支持以下几种类型：
     *  AT_SEQUENCE_NUMBER：从指定的序列号开始读取数据。
     *  AFTER_SEQUENCE_NUMBER：从指定序列号之后（不包括指定序列号）开始读取数据。
     *  TRIM_HORIZON：从最早记录（历史数据）开始读取数据。
     *  LATEST：从最新数据开始读取数据。
     *  AT_TIMESTAMP：从指定的时间戳开始读取数据。
     */
    private String getShardIterator(String shardId) {
        GetShardIteratorRequest request = new GetShardIteratorRequest()
                .withStreamName(streamName)
                .withShardId(shardId)
                .withShardIteratorType("TRIM_HORIZON"); // 从最早记录（历史数据）开始获取

        GetShardIteratorResult result = client.getShardIterator(request);
        return result.getShardIterator();
    }




    @After
    public void close() {
        // 关闭 AmazonKinesis 客户端
        client.shutdown();
    }
}
